import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import json
from collections import defaultdict
import argparse

class Compartment():
    def __init__(self, name, parameters, clear_rate, out_rate, factors, *args, **kwargs):
        self.name = name
        self.factors = factors
        self.parameters = parameters
        self.clear_rate = clear_rate
        self.out_rate = out_rate

    def __getattr__(self, item):
        # 获取参数的值
        if item in self.parameters.keys():
            return self.parameters[item]
        # 获取因子名
        if item in self.factors:
            return f'{self.name}.{item}'

    def to_pandas(self):
        return pd.DataFrame(self.output, index=self.ts, columns=self.factors)

    def vector_todict(self, y):
        assert len(y) == len(self.factors)
        return dict(zip(self.factors, np.float64(y)))

    def dict_tovector(self, dict_y):
        assert len(dict_y.keys())
        return np.array([dict_y.get(k, 0) for k in self.factors])


class Simple(Compartment):
    def __init__(self,workspace='human',config=None) -> None:
        self.factors = ['Center.drug', 'Tumor.drug', 'Peri.drug', 'Center.eaa', 'Tumor.eaa', 'Tumor.taa',
                        'Center.dimer_eaa', 'Tumor.dimer_eaa', 'Tumor.dimer_taa', 'Tumor.trimer']
        self.parameters = {}
        self.workspace = workspace
        os.makedirs(self.workspace,exist_ok=True)
        if config is not None:
            self.parameters = self.load_paramters(config)
            self.parameters['Tumor_volumn'] = 3.14 * 4 * (self.Rtumor ** 3) / 3  # L
            self.parameters['t1/2(hour)'] = np.log(2) / self.kel
            self.parameters['kdeaa'] = self.koffeaa/self.koneaa
            self.parameters['kdtaa'] = self.kofftaa/self.kontaa
        else:
            self.parameters['koneaa'] = 2e+5 / 1e+9 * 3600  # 1/nm/h
            self.parameters['koffeaa'] = 1.5e-4 * 3600  # 1/h
            self.parameters['kontaa'] = 2.94e+5 / 1e+9 * 3600  # 1/nm/h
            self.parameters['kofftaa'] = 1.38e-4 * 3600  # 1/h
            self.parameters['Center_volumn'] = 40.2e-3  # L/kg
            self.parameters['Peri_volumn'] = 211e-3  # L/kg
            #        self.parameters['CL'] = 4.61e-3 # L/h/kg
            self.parameters['CL'] = 0.0007802  # L/h/kg
            self.parameters['CLd'] = 25.2e-3  # L/h/kg
            # self.parameters['Center.Tcell'] = 0 # cells/L
            self.parameters['eaa_per_cell'] = 1e+5  # molecular / cell
            self.parameters['taa_per_cell'] = 28706  # molecular / cell
            self.parameters['P'] = 334e-5 / 24  # dm/h
            self.parameters['D'] = 0.022e-2 / 24  # dm^2/h
            self.parameters['Rcap'] = 8e-5  # dm
            self.parameters['Rkrog'] = 75e-5  # dm
            self.parameters['epsilon'] = 0.24
            self.parameters['Rtumor'] = 1e-1  # dm
            self.parameters['T'] = 10000  # h
            self.parameters['ed'] = 200  # h
            self.parameters['interval'] = 0.001  # h
            self.parameters['dose'] = 1  # nm/kg
            self.parameters['weight'] = 70  # kg
            self.parameters['tumorperg'] = 1e+8 # /g
            self.parameters['ecellperg'] = 6.49e+5 #/g
            self.update()

    def update(self):
        self.parameters['Center_volumn'] *= self.weight
        self.parameters['Peri_volumn'] *= self.weight
        self.parameters['CL'] *= self.weight
        self.parameters['CLd'] *= self.weight
        self.parameters['kel'] = self.CL / self.Center_volumn
        self.parameters['k12'] = self.CLd / self.Center_volumn
        self.parameters['k21'] = self.CLd / self.Peri_volumn
        self.parameters['dose'] *= self.weight
        self.save_parameter()
        self.parameters['Tumor_volumn'] = 3.14 * 4 * (self.Rtumor ** 3) / 3  # L
        self.parameters['t1/2(hour)'] = np.log(2) / self.kel
        self.parameters['kdeaa'] = self.koffeaa/self.koneaa
        self.parameters['kdtaa'] = self.kofftaa/self.kontaa
        

    def save_parameter(self):
        with open(f"{self.workspace}/parameters.json", 'w') as f:
            json.dump(self.parameters, f, indent=2)

    def delta(self, y0):
        delta = {f: 0 for f in y0.keys()}
        C1 = y0['Center.drug']
        C2 = y0['Peri.drug']
        C3 = y0['Tumor.drug']
        taa = y0['Tumor.taa']
        eaa = y0['Tumor.eaa']
        dimer_taa = y0['Tumor.dimer_taa']
        dimer_eaa = y0['Tumor.dimer_eaa']
        trimer = y0['Tumor.trimer']
        TD = (2 * self.P * self.Rcap / (self.Rkrog ** 2) + 6 * self.D / (self.Rtumor ** 2)) * (
                    C1 / self.Center_volumn - C3 / self.epsilon / self.Tumor_volumn)
        delta['Center.drug'] += - self.kel * C1  # 清除速率
        delta['Center.drug'] += - self.k12 * C1  # Center ->Peri
        delta['Center.drug'] += -TD * self.Tumor_volumn  # Center ->Tumor
        delta['Center.drug'] += self.k21 * C2  # Peri->Center
        delta['Peri.drug'] += self.k12 * C1  # Center->peri
        delta['Peri.drug'] += - self.k21 * C2  # Peri->center
        delta['Tumor.drug'] += TD * self.Tumor_volumn  # Center -> Tumor
        delta[
            'Tumor.drug'] += -self.koneaa * C3 * eaa / self.Tumor_volumn / self.epsilon - self.kontaa * taa * C3 / self.Tumor_volumn / self.epsilon + self.koffeaa * dimer_eaa + self.kofftaa * dimer_taa
        delta[
            'Tumor.taa'] += -self.kontaa * taa * C3 / self.Tumor_volumn / self.epsilon + self.kofftaa * dimer_taa + self.kofftaa * trimer - self.kontaa * taa * dimer_eaa / self.epsilon / self.Tumor_volumn
        delta[
            'Tumor.eaa'] += -self.koneaa * eaa * C3 / self.Tumor_volumn / self.epsilon + self.koffeaa * dimer_eaa + self.koffeaa * trimer - self.koneaa * eaa * dimer_taa / self.epsilon / self.Tumor_volumn
        delta[
            'Tumor.dimer_taa'] += self.kontaa * taa * C3 / self.Tumor_volumn / self.epsilon - self.kofftaa * dimer_taa + self.koffeaa * trimer - self.koneaa * dimer_taa * eaa / self.Tumor_volumn / self.epsilon
        delta[
            'Tumor.dimer_eaa'] += self.koneaa * eaa * C3 / self.Tumor_volumn / self.epsilon - self.koffeaa * dimer_eaa + self.kofftaa * trimer - self.kontaa * dimer_eaa * taa / self.Tumor_volumn / self.epsilon
        delta[
            'Tumor.trimer'] += self.kontaa * dimer_eaa * taa / self.Tumor_volumn / self.epsilon + self.koneaa * dimer_taa * eaa / self.Tumor_volumn / self.epsilon - self.koffeaa * trimer - self.kofftaa * trimer
        return delta

    def circleDo(self, t):
        '''
        x: 药物剂量
        t: 当前时间
        T: 给药周期
        ed: 停止给药时间
        '''
        return {'Center.drug': self.dose / self.interval if t % self.T == 0 else 0}

    def delta2(self, y0, t, *args, **kwargs):
        delta = np.zeros(len(self.factors), dtype=np.float64)
        if not isinstance(y0,dict):
            y0 = self.vector_todict(y0)
        delta += self.dict_tovector(self.delta(y0))
        delta += self.dict_tovector(self.circleDo(t))
        return delta

    def simulation(self, y0, *args, **kwargs):
        self.ts = np.arange(0, self.ed + self.interval, self.interval)
        output = []
        if isinstance(y0,dict):
            y0 = self.dict_tovector(y0)
        y0 = np.array(y0, dtype=np.float64)
        output.append(y0)
        for t in self.ts[0:-1]:
            delta = self.delta2(y0, t)
            y0 = y0 + self.interval * delta
            output.append(y0)
        self.output = np.array(output)

    def load_paramters(self,jsonfile):
        with open(jsonfile,'r') as f:
            return json.load(f)
        
    def init_condition(self):
        y0 = dict(zip(self.factors,[0 for i in self.factors]))
        y0['Tumor.taa'] = self.tumorperg * self.Tumor_volumn * self.taa_per_cell/6.023e+14
        y0['Tumor.eaa'] = self.ecellperg * self.Tumor_volumn * self.eaa_per_cell / 6.023e+14
        return y0

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config','-c',type=str,default=None, help='parameter file')
    parser.add_argument('--init','-i',type=str,default=None,help='init condition file')
    parser.add_argument('--workspace','-w',type=str,default=None,help='init condition file')
    args = parser.parse_args()
    config = 'bsabsdefault'
    if args.config:
        config = args.config
    workspace = config.rstrip('.json')
    if args.workspace:
        workspace = args.workspace
    run = Simple(workspace)
    if config.endswith('.json'):
        run.load_paramters(config)
    init_file = os.path.join(run.workspace,'init_condition.json')
    if args.init:
        init_file = args.init        
    if not os.path.exists(init_file):
        with open(init_file,'w') as f:
            json.dump(run.init_condition(),f)
        print(f'init condition file is not find!(init_file)')
        print('loading default init value')
        # raise ValueError(f'init condition file is not find!(init_file)')
    with open(init_file, 'r') as f:
        y0 = json.load(f)
    run.simulation(y0)
    data = run.to_pandas()
    data.to_csv(os.path.join(run.workspace,'data.csv'))